import cv2
import numpy as np
from matplotlib import pyplot as plt

global_predictor = None
from segment_anything import sam_model_registry, SamPredictor


def get_predictor():
    global global_predictor
    if global_predictor is not None:
        return global_predictor

    # sam_checkpoint = "/LSEM/user/zhoucan/unet/sam_vit_h_4b8939.pth"
    sam_checkpoint = "C:\\Users\\crxc\\PycharmProjects\\segment-anything\\notebooks\\sam_vit_h_4b8939.pth"
    model_type = "vit_h"

    device = "cuda"

    sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
    sam.to(device=device)

    global_predictor = SamPredictor(sam)
    return global_predictor


def predict_mask(img, points):
    predictor = get_predictor()
    img_np = img.cpu().numpy().squeeze() * 255.0
    img_np = img_np.astype(np.uint8)
    if img_np.ndim == 3:
        img_np = img_np.transpose(1, 2, 0)  # CHW to HWC
    img_normal = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
    # img_normal = img.cpu().numpy().astype(np.uint8).transpose(1, 2, 0)
    plt.imshow(img_normal)
    plt.axis('off')  # 关闭坐标轴
    plt.savefig('test.png', bbox_inches='tight', pad_inches=0)  # 保存图像
    predictor.set_image(img_normal)
    # 获取图像尺寸
    _, height, width = img.shape
    # 提取坐标和标签
    relative_coords = points[:, :2].numpy()
    labels = points[:, 2].numpy()
    # 将标签二进制化
    input_label = (labels > 0.5).astype(int).reshape(-1, )
    # 将相对坐标转换为绝对像素坐标
    absolute_coords = (relative_coords * np.array([width, height])).astype(int)

    masks, scores, logits = predictor.predict(
        point_coords=absolute_coords,
        point_labels=input_label,
        multimask_output=False,
    )
    return masks[0], scores[0]
